import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import dctn, idctn, dstn, idstn

# 参数设置
Nx, Ny = 32, 32  # 示例中使用 Nx 和 Ny 的值

# 网格间距
dx = 2 / Nx
dy = 2 / Ny

# 创建网格
x = np.linspace(-1, 1, Nx, endpoint=False)
y = np.linspace(-1, 1, Ny, endpoint=False)
X, Y = np.meshgrid(x, y, indexing='ij')

# 精确解
u_analytic = np.cos(np.pi * X) * np.cos(np.pi * Y)

# 修改波数
kx = np.fft.fftfreq(Nx, d=dx) * 2 * np.pi
ky = np.fft.fftfreq(Ny, d=dy) * 2 * np.pi

# 2D DCT of f (Right hand side)
f = -np.pi**2 * u_analytic
fhat = dctn(f, type=2)

# 构造 MWX 和 MWY 的网格
KX, KY = np.meshgrid(kx, ky, indexing='ij')
MWX = -KX**2 / dx**2
MWY = -KY**2 / dy**2

# 解方程
uhat = fhat / (MWX + MWY)
uhat[0, 0] = 0  # 由于 uhat(0,0) 可能未定义，我们将其设置为 0

# 逆 2D DCT
u = idctn(uhat, type=2)
u_numeric = np.real(u)  # 使用数值解的实部

# 计算误差
error = np.abs(u_numeric - u_analytic)

# 绘制数值解、解析解和误差
plt.figure(figsize=(15, 5))

# 数值解
plt.subplot(1, 3, 1)
plt.title("Numerical Solution")
plt.imshow(u_numeric, extent=[-1, 1, -1, 1], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

# 精确解
plt.subplot(1, 3, 2)
plt.title("Analytical Solution")
plt.imshow(u_analytic, extent=[-1, 1, -1, 1], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

# 误差
plt.subplot(1, 3, 3)
plt.title("Error between Numerical and Analytical Solutions")
plt.imshow(error, extent=[-1, 1, -1, 1], origin='lower', cmap='coolwarm')
plt.colorbar(label='Error')

plt.tight_layout()
plt.show()